#ifndef	gemat_included
#define	gemat_included

__global__	void	_gematfp(uint32_t	H,	uint32_t	C,	half	*att) {
    uint32_t	id=blockIdx.x*blockDim.x+threadIdx.x,	c=id%C,	s=min((((c+1)>>2)+1)<<2,C),i;
    half	*a=att+id*C;
    half2	ma= {a[c],a[c]};
    for(i=0;	i+1<=c;	i+=2) ma=__hmax2(ma,*(half2*)(a+i));
    ma.y=ma.x=__hmax(ma.x,ma.y);
    for(i=0;	i+1<=c;	i+=2) *(half2*)(a+i)=h2exp(__hsub2(*(half2*)(a+i),ma));
    if(i<=c)	a[c]=hexp(__hsub(a[c],ma.x));
    uint16_t	*b=(uint16_t*)a;
    for(i=c+1;	i<s;	i++)	b[i]=0;
    for(i=s;	i<C;	i+=4)	*(short4*)(b+i)= {};
}

__global__	void	_gematf(uint32_t	R,	half	*u,	half	*v,	half	*out) {
    uint32_t	id=(blockIdx.x*blockDim.x+threadIdx.x)<<3,r=id%R,c=id/R;
    half	*pu=u+c*4*R+3*R+r,	*pv=v+c*R+r;
    *(half8*)(out+id)=_hmul8(*(half8*)pu,*(half8*)pv);
}

__global__	void	_gematb(uint32_t	R,	half	*u,	half	*v,	half	*gin,	half	*U) {
    uint32_t	id=(blockIdx.x*blockDim.x+threadIdx.x)<<3,r=id%R,c=id/R;
    half	*pu=u+c*4*R+3*R+r,*pv=v+c*R+r;
    half	*du=U+c*4*R+3*R+r,*dv=gin+c*R+r;
    *(half8*)du=_hmul8(*(half8*)pv,*(half8*)(gin+id));
    *(half8*)dv=_hmul8(*(half8*)pu,*(half8*)(gin+id));
}

__global__	void	_gematba(half	*gin,	half	*att) {
    uint32_t	id=(blockIdx.x*blockDim.x+threadIdx.x)<<3;
    *(half8*)(gin+id)=_hmul8(*(half8*)(att+id),*(half8*)(gin+id));
}

template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
struct	gemat {
    static	tensor<half,C*C*H>	da;
    static	tensor<half,S*H*C>	gi;
    static	tensor<half,4*S*H*C>	gx;
    tensor<half,C*C*H>	at;
    tensor<half,R*C>	in;
    tensor<half,S*H*C>	va;
    tensor<half,S*H*C>	tmp;
    tensor<half,4*S*H*C>	x0,xo;
    rmsnorm<4*S*H,C,4*H>	nx;
    rmsnorm<S*H,C,H>	nv;
    linear<R,4*S*H,C>	x;
    ema<4*S*H,C>	m;
    linear<S*H,R,C>	o;

    void	save(FILE	*F) {
        x.save(F);
        m.save(F);
        o.save(F);
    }

    void	load(FILE	*F) {
        x.load(F);
        m.load(F);
        o.load(F);
    }

    uint64_t	size(void) {
        return	x.size()+m.size()+o.size();
    }

    void	fw(tensor<half,R*C>	&inp) {
        half	alf=__float2half_rn(2/sqrtf(S)),alf1=__float2half_rn(1),bet=__float2half_rn(0);
        cudaMemcpy(in.p,inp.p,R*C*sizeof(half),cudaMemcpyDeviceToDevice);
        x.fwt(in,x0,0);
        m.fw(x0,xo);
        nx.fw(xo);
        cublasGemmStridedBatchedEx(handle,CUBLAS_OP_T,CUBLAS_OP_N,C,C,S,&alf,xo.p,CUDA_R_16F,4*S*H,S,xo.p+S*H,CUDA_R_16F,4*S*H,S,&bet,at.p,CUDA_R_16F,C,C*C,H,CUBLAS_COMPUTE_16F,CUBLAS_GEMM_DEFAULT);
        uint32_t	w0=warp_size(1,C*H);
        _gematfp<<<C*H/w0,w0>>>(H,C,at.p);
        cublasGemmStridedBatchedEx(handle,CUBLAS_OP_N,CUBLAS_OP_N,S,C,C,&alf1,xo.p+2*S*H,CUDA_R_16F,4*S*H,S,at.p,CUDA_R_16F,C,C*C,&bet,va.p,CUDA_R_16F,S*H,S,H,CUBLAS_COMPUTE_16F,CUBLAS_GEMM_DEFAULT);
        nv.fw(va);
        uint32_t	w1=warp_size(8,S*H*C);
        _gematf<<<S*H*C/8/w1,w1>>>(S*H,xo.p,va.p,tmp.p);
        o.fw(tmp,inp,resnet_scaling);
    }

    void	bk(tensor<half,R*C>	&gin) {
        half	alf=__float2half_rn(2/sqrtf(S)),alf1=__float2half_rn(1),bet=__float2half_rn(0);
        o.bk(tmp,gin,gi,0);
        uint32_t	w0=warp_size(8,S*H*C);
        _gematb<<<S*H*C/8/w0,w0>>>(S*H,xo.p,va.p,gi.p,gx.p);
        nv.bk(va,gi);
        cublasGemmStridedBatchedEx(handle,CUBLAS_OP_T,CUBLAS_OP_N,C,C,S,&alf1,xo.p+2*S*H,CUDA_R_16F,4*S*H,S,gi.p,CUDA_R_16F,S*H,S,&bet,da.p,CUDA_R_16F,C,C*C,H,CUBLAS_COMPUTE_16F,CUBLAS_GEMM_DEFAULT);
        cublasGemmStridedBatchedEx(handle,CUBLAS_OP_N,CUBLAS_OP_T,S,C,C,&alf1,gi.p,CUDA_R_16F,S*H,S,at.p,CUDA_R_16F,C,C*C,&bet,gx.p+2*S*H,CUDA_R_16F,4*S*H,S,H,CUBLAS_COMPUTE_16F,CUBLAS_GEMM_DEFAULT);
        uint32_t	w1=warp_size(8,C*C*H);
        _gematba<<<C*C*H/8/w1,w1>>>(da.p,at.p);
        cublasGemmStridedBatchedEx(handle,CUBLAS_OP_N,CUBLAS_OP_T,S,C,C,&alf,xo.p+S*H,CUDA_R_16F,4*S*H,S,da.p,CUDA_R_16F,C,C*C,&bet,gx.p,CUDA_R_16F,4*S*H,S,H,CUBLAS_COMPUTE_16F,CUBLAS_GEMM_DEFAULT);
        cublasGemmStridedBatchedEx(handle,CUBLAS_OP_N,CUBLAS_OP_N,S,C,C,&alf,xo.p,CUDA_R_16F,4*S*H,S,da.p,CUDA_R_16F,C,C*C,&bet,gx.p+S*H,CUDA_R_16F,4*S*H,S,H,CUBLAS_COMPUTE_16F,CUBLAS_GEMM_DEFAULT);
        nx.bk(xo,gx);
        m.bk(x0,gx,xo);
        x.bk(in,xo,gin,resnet_scaling);
    }
};

template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
tensor<half,C*C*H>	gemat<R,C,H,S>::da;
template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
tensor<half,S*H*C>	gemat<R,C,H,S>::gi;
template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
tensor<half,4*S*H*C>	gemat<R,C,H,S>::gx;

#endif

